package com.k12.user.interceptor;

import org.apache.ibatis.builder.StaticSqlSource;
import org.apache.ibatis.executor.Executor;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.mapping.SqlSource;
import org.apache.ibatis.plugin.*;
import org.apache.ibatis.session.ResultHandler;
import org.apache.ibatis.session.RowBounds;

import java.lang.reflect.Field;
import java.util.Properties;

@Intercepts({
        @Signature(type = Executor.class, method = "update", args = {MappedStatement.class, Object.class}),
        @Signature(type = Executor.class, method = "query", args = {MappedStatement.class, Object.class, RowBounds.class, ResultHandler.class})
})
public class MySqlRewriteInterceptor implements Interceptor {

    // 表名映射规则
    private Properties tableMappings = new Properties();
    
    // 是否启用强制索引
    private boolean forceIndex = false;
    
    // 租户ID列名
    private String tenantIdColumn;
    
    @Override
    public Object intercept(Invocation invocation) throws Throwable {
        // 获取原始SQL信息
        Object[] args = invocation.getArgs();
        MappedStatement mappedStatement = (MappedStatement) args[0];
        Object parameter = args[1];
        
        BoundSql boundSql = mappedStatement.getBoundSql(parameter);
        String originalSql = boundSql.getSql();

        // 重写SQL
        String newSql = rewriteSql(originalSql);
        
        // 使用新SQL替换原SQL
        resetSqlSource(mappedStatement, newSql);
        
        return invocation.proceed();
    }


    private void resetSqlSource(MappedStatement ms, String newSql) {
        try {
            Field sqlSourceField = MappedStatement.class.getDeclaredField("sqlSource");
            sqlSourceField.setAccessible(true);

            SqlSource originalSqlSource = (SqlSource) sqlSourceField.get(ms);
            SqlSource newSqlSource = new StaticSqlSource(
                    ms.getConfiguration(),
                    newSql,
                    originalSqlSource.getBoundSql(null).getParameterMappings()
            );

            sqlSourceField.set(ms, newSqlSource);
            sqlSourceField.setAccessible(false);
        } catch (Exception e) {
            throw new RuntimeException("Failed to modify SqlSource", e);
        }
    }
    
    private String rewriteSql(String originalSql) {
        String newSql = originalSql;

        // 1. 表名替换
        for (Object key : tableMappings.keySet()) {
            String oldTable = (String) key;
            String newTable = tableMappings.getProperty(oldTable);
            newSql = newSql.replaceAll("(?i)\\b" + oldTable + "\\b", newTable);
        }

        // 2. 添加租户过滤条件
        if (tenantIdColumn != null && !tenantIdColumn.isEmpty()) {
            newSql = addTenantCondition(newSql);
        }

        // 3. 强制索引
        if (forceIndex) {
            newSql = addForceIndex(newSql);
        }

        return newSql;
    }
    
    private String addTenantCondition(String sql) {
        // 简单实现：在WHERE条件中添加租户ID过滤
        // 实际项目需要更复杂的解析逻辑
        if (sql.toUpperCase().contains(" WHERE ")) {
            return sql.replaceAll("(?i) WHERE ", " WHERE " + tenantIdColumn + " = 4 AND ");
        } else {
            return sql + " WHERE " + tenantIdColumn + " = #{tenantId}";
        }
    }
    
    private String addForceIndex(String sql) {
        // 为SELECT语句添加FORCE INDEX
        if (sql.toUpperCase().startsWith("SELECT")) {
            return sql.replaceAll("(?i)FROM\\s+(\\w+)", "FROM $1 FORCE INDEX (PRIMARY)");
        }
        return sql;
    }
    @Override
    public Object plugin(Object target) {
        return Plugin.wrap(target, this);
    }
    
    @Override
    public void setProperties(Properties properties) {
        // 加载表名映射配置
        if (properties.containsKey("tableMappings")) {
            String mappings = properties.getProperty("tableMappings");
            String[] pairs = mappings.split(",");
            for (String pair : pairs) {
                String[] kv = pair.split("=");
                if (kv.length == 2) {
                    tableMappings.setProperty(kv[0].trim(), kv[1].trim());
                }
            }
        }
        
        // 是否启用强制索引
        if (properties.containsKey("forceIndex")) {
            this.forceIndex = Boolean.parseBoolean(properties.getProperty("forceIndex"));
        }
        
        // 租户ID列名
        if (properties.containsKey("tenantIdColumn")) {
            this.tenantIdColumn = properties.getProperty("tenantIdColumn");
        }
    }

}